# _*_ coding: utf-8 _*_
"""
Time:     2024/9/14
Author:   2051261
File:     task.py
"""
import asyncio
import logging
import traceback
from collections import defaultdict
from typing import Dict, List

from extractor.flow.node_model import node_name_flow
from extractor.flow.task_model import TaskModel
from extractor.flow.node_model.log_node import RecordLog

logger = logging.getLogger(__name__)

sem = asyncio.Semaphore(30)  # 最大并发数为3


class Cache:
    def __init__(self):
        self.value = None
        self.key = None
        self.table_columns = None

    async def initialize(self, key, value):
        if self.value is None:
            self.key = key
            self.value = value

    async def update(self, value):
        self.value = self.value if self.value > value else value

    async def get_key(self):
        return self.key

    async def get_value(self):
        return self.value

    async def cache_table_columns(self, table_name: str, columns: list[str]):
        if self.table_columns is None:
            self.table_columns = {f"{table_name}": columns}
        else:
            self.table_columns.update({f"{table_name}": columns})

    async def get_table_columns(self) -> Dict[str, List[str]]:
        return self.table_columns


class TaskEntry:

    def __init__(self, task_model: TaskModel):
        self.task_model = task_model
        self.cache = Cache()   # 每个任务实例化一个对象时会使用一个cache实例来存储每次运行时的缓存，避免每次查库

    async def run(self):
        async with sem:
            context = {}
            record = defaultdict(int)
            record_log = RecordLog(**self.task_model.log)
            output = {"code": 0}
            stat_info = {}
            try:
                for node in self.task_model["flow"]:
                    node_name = node["node_name"]
                    instance = node_name_flow[node_name]()
                    context = await instance(record_log, context, self.cache, record, node["kwargs"])
            except Exception as e:
                stat_info.update({"error": str(e)[:100]})
                output["code"] = 1
                logger.error(f'Task: {self.task_model.task_name} execution failed, error: {e} \n'
                             f' {traceback.format_exc()}')
            finally:
                stat_info.update(record)
                await record_log.end(output, stat_info)
